import math

from matplotlib import pyplot
import numpy

from fitter import FitException
from bfitter import BinnedFitter

class BinnedFitterGauss(BinnedFitter):
    fcn_latex = r'$A \exp\left[-\frac{(x - \mu)^2}{\sigma^2}\right]$'
    fcn_str="{0}*math.exp(-0.5*(((x-{1})/{2})**2))"

    def __init__(self, data, par_names=r'$A$ $\mu$ $\sigma$'.split()):
        super(BinnedFitterGauss, self).__init__(data, par_names, self.fcn_str)
        self._par_text_offset = 0.5

    def _set_initpars(self, frange):
        if not self.hist:
            raise FitException("BUG in gauss.py. Need to set self.hist before calling _set_initpars.")

        bins = self.hist.bins()
        
        max_bin = max(bins)

        max_bin_index = 0
        for i in xrange(len(bins)):
            if bins[i] == max_bin:
                max_bin_index = i
                break
        
        return [max_bin, self.data_mean, self.data_std]
                                                             



class BinnedFitterTwoGauss(BinnedFitter):
    fcn_str   = BinnedFitterGauss.fcn_str + " + {3}*math.exp(-0.5*(((x-{4})/{5})**2))"
    fcn_latex = '$A_1e^{-(x - \\mu_1)^2 / \\sigma_1^2} + A_2e^{-(x - \\mu_2)^2 / \\sigma_2^2}$'

    def __init__(self, hist, par_names="$A_1$ $\\mu_1$ $\\sigma_1$ $A_2$ $\\mu_2$ $\\sigma_2$".split()):
        super(BinnedFitterTwoGauss, self).__init__(hist, par_names=par_names, fcn_str=self.fcn_str)
        
        self._par_text_offset = 0.35

        
    
class BinnedFitterGaussBG(BinnedFitterGauss):
    fcn_latex = "$A e^{-(x - \\mu)^2 / \\sigma^2}$ \n $ + p_0 + p_1x + p_2x^2$"
    
    def __init__(self, data, par_names='A $\\mu$ $\\sigma$ $p_0$ $p_1$ $p_2$'.split()):
        super(BinnedFitterGaussBG, self).__init__(data, par_names)
        self._par_text_offset = 0.325

    def fcn(self, p, x):
        return p[0]*math.exp(-0.5*(((x-p[1])/p[2])**2)) + p[3] + p[4]*x + p[5]*x**2

    def _set_initpars(self, frange):
        if not self.hist:
            raise FitException("BUG in gauss.py. Need to set self.hist before calling _set_initpars.")

        bins = self.hist.bins()
        
        max_bin = max(bins)

        max_bin_index = 0
        for i in xrange(len(bins)):
            if bins[i] == max_bin:
                max_bin_index = i
                break
        
        return [max_bin, self.data_mean, self.data_std, 1, 1, 1]

        
